import whisper
import opencc
import sys
import numpy as np

def transcribe(audio_file):
    # 加载 Whisper 模型
    model = whisper.load_model("base")
    
    # 转录音频
    result = model.transcribe(audio_file, word_timestamps=True)
    # print(result)
    
    # 初始化 opencc 转换器，转换繁体到简体
    cc = opencc.OpenCC('t2s.json')  # 繁体转简体，t2s.json 是转换配置文件
    
    result['text'] = cc.convert(result['text'])
    # 获取时间信息
    segments = result['segments']  # 获取每个段落的时间戳信息
    
    for segment in segments:
        segment['start'] = f"{np.round(segment['start'], 2)}"
        segment['end'] = f"{np.round(segment['end'], 2)}"
        segment['text'] = cc.convert(segment['text'])
        del segment['tokens']
        del segment['words']
        del segment['temperature']
        del segment['avg_logprob']
        del segment['compression_ratio']
        del segment['no_speech_prob']
    
    return result

if __name__ == "__main__":
    audio_file = sys.argv[1]  # 获取命令行传入的音频文件路径
    result = transcribe(audio_file)

    result['code'] = 0
    
    # 打印转录文本（简体）
    # print("Transcribed Text (Simplified):")
    print(result)


